"""
代码格式化工具

这个模块提供了用于格式化文本中代码块的工具类。
支持将文本中的代码块转换为带有语法高亮的格式。
"""
from typing import List, Tuple
import re
from prompt_toolkit.formatted_text import HTML
from pygments import highlight
from pygments.lexers import get_lexer_by_name, TextLexer
from pygments.formatters import Terminal256Formatter  # 修改为终端格式化器
from pygments.util import ClassNotFound
from pygments.token import Token

class CodeFormatter:
    """
    代码格式化工具类
    
    用于将文本中的代码块转换为带有语法高亮的格式。
    支持多种编程语言的代码高亮。
    """
    
    @staticmethod
    def _extract_code_blocks(text: str) -> List[Tuple[str, str, str]]:
        """
        从文本中提取所有代码块
        
        Args:
            text: 包含代码块的文本
            
        Returns:
            List[Tuple[str, str, str]]: (前缀文本, 语言, 代码内容)的列表
            
        示例:
            >>> text = "Hello\\n```python\\nprint('hi')\\n```\\nWorld"
            >>> CodeFormatter._extract_code_blocks(text)
            [('Hello\\n', 'python', "print('hi')")]
        """
        # 分割文本和代码块的正则表达式[\\n]+
        pattern = r'```([\w]*)(.*?)```'
        
        # 查找所有匹配项
        matches = list(re.finditer(pattern, text, re.DOTALL))
        blocks = []
        
        last_end = 0
        for match in matches:
            # 提取匹配内容
            prefix = text[last_end:match.start()]
            lang = match.group(1).strip()
            code = match.group(2).rstrip()
            blocks.append((prefix, lang, code))
            last_end = match.end()
        
        # 添加最后一段文本（如果有）
        if last_end < len(text):
            blocks.append((text[last_end:], '', ''))
            
        return blocks
    
    @staticmethod
    def _highlight_code(code: str, language: str = '') -> str:
        """
        对代码应用语法高亮
        
        Args:
            code: 要高亮的代码
            language: 编程语言名称
            
        Returns:
            str: 带ANSI颜色代码的文本
        """
        try:
            if language:
                lexer = get_lexer_by_name(language, stripall=True)
            else:
                lexer = TextLexer()
            
            # 使用Terminal256Formatter，设置为白色背景主题
            formatter = Terminal256Formatter(
                style='default',  # 使用默认风格
                bg='light'       # 使用亮色背景
            )
            
            # 获取高亮后的代码
            highlighted = highlight(code, lexer, formatter)
            
            # 处理代码块
            lines = highlighted.rstrip().split('\n')
            
            # 计算实际显示宽度
            ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
            clean_lines = [ansi_escape.sub('', line) for line in lines]
            max_length = max(len(line.rstrip()) for line in clean_lines)
            
            # 添加上边框
            result = [f"\033[48;5;255m\033[38;5;240m┌{'─' * (max_length + 2)}┐\033[0m"]
            
            # 添加代码内容
            for line in lines:
                # 获取当前行实际显示宽度
                clean_line = ansi_escape.sub('', line.rstrip())
                padding = max_length - len(clean_line)
                
                # 构建格式化行
                formatted_line = (
                    f"\033[48;5;255m\033[38;5;240m│ "  # 左边框
                    f"\033[48;5;255m{line.rstrip()}"    # 代码内容
                    f"\033[48;5;255m{' ' * padding}"    # 填充空格
                    f"\033[48;5;255m\033[38;5;240m │"   # 右边框
                    f"\033[0m"                          # 重置样式
                )
                result.append(formatted_line)
            
            # 添加下边框
            result.append(f"\033[48;5;255m\033[38;5;240m└{'─' * (max_length + 2)}┐\033[0m")
            
            return '\n'.join(result)
            
        except ClassNotFound:
            # 如果找不到指定语言的lexer，使用普通文本格式
            lines = code.rstrip().split('\n')
            max_length = max(len(line.rstrip()) for line in lines)
            
            result = [f"\033[48;5;255m\033[38;5;240m┌{'─' * (max_length + 2)}┐\033[0m"]
            
            for line in lines:
                clean_line = line.rstrip()
                padding = max_length - len(clean_line)
                formatted_line = (
                    f"\033[48;5;255m\033[38;5;240m│ "  # 左边框
                    f"\033[48;5;255m\033[38;5;16m{clean_line}"  # 文本内容（黑色）
                    f"\033[48;5;255m{' ' * padding}"  # 填充空格
                    f"\033[48;5;255m\033[38;5;240m │\033[0m"  # 右边框
                )
                result.append(formatted_line)
            
            result.append(f"\033[48;5;255m\033[38;5;240m└{'─' * (max_length + 2)}┐\033[0m")
            
            return '\n'.join(result)
    
    @staticmethod
    def format_text(text: str) -> str:  # 返回类型改为str
        """
        格式化包含代码块的文本
        
        将文本中的代码块转换为带有ANSI颜色代码的格式。
        非代码部分保持原样。
        
        Args:
            text: 要格式化的文本
            
        Returns:
            str: 格式化后的文本
        """
        # 提取所有代码块
        blocks = CodeFormatter._extract_code_blocks(text)
        
        # 构建结果
        result = []
        
        for prefix, lang, code in blocks:
            if prefix:
                # 添加普通文本
                result.append(prefix)
            
            if code:
                # 添加高亮的代码块
                highlighted = CodeFormatter._highlight_code(code, lang)
                result.append(highlighted)
        
        # 将所有部分连接成字符串
        return ''.join(result) 